#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
布局保持模块
根据OCR结果的坐标信息进行原样排版
"""

import json
import os

class LayoutPreserver:
    def __init__(self):
        pass
    
    def preserve_layout_from_json(self, json_file_path, output_format="text"):
        """
        根据OCR JSON结果文件进行原样排版
        
        Args:
            json_file_path (str): OCR结果JSON文件路径
            output_format (str): 输出格式 ("text", "html", "markdown")
            
        Returns:
            str: 保持原布局的文本内容
        """
        try:
            # 读取JSON文件
            with open(json_file_path, 'r', encoding='utf-8') as f:
                ocr_data = json.load(f)
            
            # 提取文本和坐标信息
            texts = ocr_data.get("rec_texts", [])
            boxes = ocr_data.get("rec_boxes", [])
            
            if not texts or not boxes:
                return "未找到有效的OCR结果数据"
            
            # 根据坐标进行排版
            if output_format == "text":
                return self._layout_as_text(texts, boxes)
            elif output_format == "html":
                return self._layout_as_html(texts, boxes)
            elif output_format == "markdown":
                return self._layout_as_markdown(texts, boxes)
            else:
                return self._layout_as_text(texts, boxes)
                
        except Exception as e:
            return f"处理OCR结果时发生错误: {str(e)}"
    
    def preserve_layout_from_page_results(self, page_results, page_index=1, output_format="text"):
        """
        根据单页OCR结果进行原样排版
        
        Args:
            page_results (list): 单页OCR结果列表
            page_index (int): 页码索引
            output_format (str): 输出格式 ("text", "html", "markdown")
            
        Returns:
            str: 保持原布局的文本内容
        """
        try:
            # 提取文本和坐标信息
            texts = [item["text"] for item in page_results]
            boxes = [item["coordinates"] for item in page_results]
            
            if not texts or not boxes:
                return "未找到有效的OCR结果数据"
            
            # 根据坐标进行排版
            if output_format == "text":
                return self._layout_as_text(texts, boxes)
            elif output_format == "html":
                return self._layout_as_html_single_page(texts, boxes, page_index)
            elif output_format == "markdown":
                return self._layout_as_markdown(texts, boxes)
            else:
                return self._layout_as_text(texts, boxes)
                
        except Exception as e:
            return f"处理OCR结果时发生错误: {str(e)}"
    
    def preserve_layout_from_multi_page_results(self, all_page_results, output_format="text"):
        """
        根据多页OCR结果进行原样排版
        
        Args:
            all_page_results (dict): 所有页面的OCR结果，key为页面路径，value为结果列表
            output_format (str): 输出格式 ("text", "html", "markdown")
            
        Returns:
            str: 保持原布局的文本内容
        """
        try:
            if output_format == "html":
                return self._layout_as_html_multi_page(all_page_results)
            else:
                # 对于文本和其他格式，按页面顺序连接
                result_parts = []
                for i, (page_path, page_results) in enumerate(all_page_results.items()):
                    page_content = self.preserve_layout_from_page_results(page_results, i+1, output_format)
                    if output_format == "text":
                        result_parts.append(f"========== 页面 {i+1} ==========\n{page_content}")
                    else:
                        result_parts.append(page_content)
                return "\n\n".join(result_parts)
                
        except Exception as e:
            return f"处理多页OCR结果时发生错误: {str(e)}"
    
    def preserve_layout_from_detailed_results(self, detailed_results, output_format="text"):
        """
        根据详细的OCR结果进行原样排版
        
        Args:
            detailed_results (list): 详细的OCR结果列表
            output_format (str): 输出格式 ("text", "html", "markdown")
            
        Returns:
            str: 保持原布局的文本内容
        """
        try:
            # 提取文本和坐标信息
            texts = [item["text"] for item in detailed_results]
            boxes = [item["coordinates"] for item in detailed_results]
            
            if not texts or not boxes:
                return "未找到有效的OCR结果数据"
            
            # 根据坐标进行排版
            if output_format == "text":
                return self._layout_as_text(texts, boxes)
            elif output_format == "html":
                return self._layout_as_html(texts, boxes)
            elif output_format == "markdown":
                return self._layout_as_markdown(texts, boxes)
            else:
                return self._layout_as_text(texts, boxes)
                
        except Exception as e:
            return f"处理OCR结果时发生错误: {str(e)}"
    
    def _layout_as_text(self, texts, boxes):
        """
        将OCR结果按坐标布局为文本格式
        
        Args:
            texts (list): 文本列表
            boxes (list): 坐标框列表
            
        Returns:
            str: 保持原布局的文本
        """
        if not texts or not boxes:
            return ""
        
        # 计算页面尺寸（基于坐标）
        max_x = max([box[2] for box in boxes]) if boxes else 0
        max_y = max([box[3] for box in boxes]) if boxes else 0
        
        # 创建二维网格表示页面
        # 为了简化处理，我们按行分组
        lines = self._group_by_lines(boxes, texts)
        
        # 构建结果文本
        result_lines = []
        for line in lines:
            line_text = "".join([item["text"] for item in line])
            result_lines.append(line_text)
        
        return "\n".join(result_lines)
    
    def _layout_as_html_single_page(self, texts, boxes, page_index=1):
        """
        将单页OCR结果按坐标布局为HTML格式
        
        Args:
            texts (list): 文本列表
            boxes (list): 坐标框列表
            page_index (int): 页码索引
            
        Returns:
            str: 保持原布局的HTML
        """
        if not texts or not boxes:
            return f'<div class="page" id="page-{page_index}"></div>'
        
        # 计算页面尺寸
        max_x = max([box[2] for box in boxes]) if boxes else 800
        max_y = max([box[3] for box in boxes]) if boxes else 600
        
        # 分组为行
        lines = self._group_by_lines(boxes, texts)
        
        # 构建HTML
        html_parts = [f'<div class="page" id="page-{page_index}" style="position: relative; width: {max_x}px; height: {max_y}px; border: 1px solid #ccc; margin-bottom: 20px;">']
        
        for line in lines:
            for item in line:
                x, y, _, _ = item["box"]
                text = item["text"]
                # 简化处理，实际应用中可能需要更精确的定位
                html_parts.append(f'<div class="ocr-text" style="position: absolute; left: {x}px; top: {y}px; white-space: nowrap;">{text}</div>')
        
        html_parts.append('</div>')
        return "".join(html_parts)
    
    def _layout_as_html_multi_page(self, all_page_results):
        """
        将多页OCR结果按坐标布局为HTML格式
        
        Args:
            all_page_results (dict): 所有页面的OCR结果
            
        Returns:
            str: 保持原布局的HTML
        """
        # 构建HTML
        html_parts = ['<html><head><meta charset="UTF-8"><style>']
        html_parts.append('body { font-family: sans-serif; }')
        html_parts.append('.page { position: relative; border: 1px solid #ccc; margin-bottom: 20px; }')
        html_parts.append('.ocr-text { position: absolute; white-space: nowrap; }')
        html_parts.append('</style></head><body>')
        
        # 按页面处理
        for i, (page_path, page_results) in enumerate(all_page_results.items()):
            page_content = self.preserve_layout_from_page_results(page_results, i+1, "html")
            html_parts.append(page_content)
        
        html_parts.append('</body></html>')
        return "".join(html_parts)
    
    def _layout_as_html(self, texts, boxes):
        """
        将OCR结果按坐标布局为HTML格式（向后兼容）
        
        Args:
            texts (list): 文本列表
            boxes (list): 坐标框列表
            
        Returns:
            str: 保持原布局的HTML
        """
        if not texts or not boxes:
            return "<html><body></body></html>"
        
        # 计算页面尺寸
        max_x = max([box[2] for box in boxes]) if boxes else 800
        max_y = max([box[3] for box in boxes]) if boxes else 600
        
        # 分组为行
        lines = self._group_by_lines(boxes, texts)
        
        # 构建HTML
        html_parts = ['<html><head><meta charset="UTF-8"><style>']
        html_parts.append('body { font-family: sans-serif; position: relative; }')
        html_parts.append('.ocr-text { position: absolute; white-space: nowrap; }')
        html_parts.append('</style></head><body>')
        
        for line in lines:
            for item in line:
                x, y, _, _ = item["box"]
                text = item["text"]
                # 简化处理，实际应用中可能需要更精确的定位
                html_parts.append(f'<div class="ocr-text" style="left: {x}px; top: {y}px;">{text}</div>')
        
        html_parts.append('</body></html>')
        return "".join(html_parts)
    
    def _layout_as_markdown(self, texts, boxes):
        """
        将OCR结果按坐标布局为Markdown格式
        
        Args:
            texts (list): 文本列表
            boxes (list): 坐标框列表
            
        Returns:
            str: 保持原布局的Markdown
        """
        # Markdown格式难以精确保持原布局，这里返回简单的文本格式
        text_layout = self._layout_as_text(texts, boxes)
        return text_layout
    
    def _group_by_lines(self, boxes, texts):
        """
        根据Y坐标将文本分组为行
        
        Args:
            boxes (list): 坐标框列表
            texts (list): 文本列表
            
        Returns:
            list: 按行分组的结果
        """
        if not boxes or not texts:
            return []
        
        # 创建文本和坐标项列表
        items = []
        for i, (text, box) in enumerate(zip(texts, boxes)):
            items.append({
                "text": text,
                "box": box,
                "index": i
            })
        
        # 按Y坐标排序
        items.sort(key=lambda x: x["box"][1])  # 按top坐标排序
        
        # 分组为行（Y坐标相近的为同一行）
        lines = []
        current_line = []
        current_y = None
        line_threshold = 10  # 同一行Y坐标的最大差异
        
        for item in items:
            y = item["box"][1]
            
            if current_y is None or abs(y - current_y) <= line_threshold:
                # 同一行
                current_line.append(item)
                if current_y is None:
                    current_y = y
            else:
                # 新行
                if current_line:
                    # 按X坐标排序
                    current_line.sort(key=lambda x: x["box"][0])
                    lines.append(current_line)
                current_line = [item]
                current_y = y
        
        # 添加最后一行
        if current_line:
            current_line.sort(key=lambda x: x["box"][0])
            lines.append(current_line)
        
        return lines

def main():
    """测试布局保持功能"""
    layout_preserver = LayoutPreserver()
    
    # 测试JSON文件处理
    json_file = os.path.join("ocr", "output", "1_res.json")
    if os.path.exists(json_file):
        print("=== 原样排版结果 ===")
        result = layout_preserver.preserve_layout_from_json(json_file)
        print(result)
        print("\n=== HTML格式 ===")
        html_result = layout_preserver.preserve_layout_from_json(json_file, "html")
        print(html_result[:200] + "..." if len(html_result) > 200 else html_result)
    else:
        print(f"测试文件 {json_file} 不存在")

if __name__ == "__main__":
    main()